//
// See the file COPYING for redistribution information.
//

#include <config.h>
#include <iostream>             // for std::cerr
#include <map>
#include <time.h>
#include <fcntl.h>
#include <errno.h>
#include <Naipax/Common.hpp>
#include <Naipax/Error.hpp>
#include <Naipax/Http.hpp>
#include <Naipax/Connection.hpp>

using namespace std;
using namespace Naipax;

ConnectionBuffer::ConnectionBuffer()
        : mStatus(STATUS_INIT), mFd(-1), mGotStatusCode(false),
          mExpireTime((time_t) -1) {
}

void ConnectionBuffer::Reset() {
    if (mFd != -1) {
        if (close(mFd) != 0) {
            cerr << strerror(errno) << " in function "
                 << __FUNCTION__ << endl;
        }
    }
    mStatus = ConnectionBuffer::STATUS_INIT;

    mFd = -1;
    
    mUrl = "";
    mWriteBuf = "";
    mWriteCurPos = 0;
    mReadBuf = "";
    mExpireTime = (time_t) -1;
    mGotStatusCode = false;
}

void ConnectionBuffer::WriteData() {
    ssize_t WriteRes = write(mFd, mWriteBuf.c_str() + mWriteCurPos,
                             mWriteBuf.length() - mWriteCurPos);
    if (WriteRes < 0) {
        if (errno != EAGAIN) {
            cerr << "Socket write error: " << strerror(errno)
                 << " in function " << __FUNCTION__ << endl;
            Reset();
            NAIPAX_THROW(Error::CANNOT_WRITE);
        }
    }
    else {
        mWriteCurPos += (uint32_t) WriteRes;
    }
                        
    if (mWriteCurPos == mWriteBuf.length()) {
        // If there is no more data to be written, then
        // half-close the socket
        if (shutdown(mFd, SHUT_WR) != 0) {
            cerr << strerror(errno)
                 << " in function " << __FUNCTION__ << endl;
            NAIPAX_THROW(Error::CANNOT_CLOSE_SOCKET);
        }
        mStatus = ConnectionBuffer::STATUS_WRITTEN;
    }
}

void ConnectionBuffer::ReadData() {
    static const uint32_t sBufSize = 1048576;
    char Buf[sBufSize];
    ssize_t ReadRes = read(mFd, Buf, (size_t) sBufSize);
    if (ReadRes < 0) {
        if (errno != EAGAIN) {
            cerr << "Socket read error: "
                 << strerror(errno)
                 << " in function " << __FUNCTION__ << endl;
            Reset();
            NAIPAX_THROW(Error::CANNOT_READ);
        }
    }
    else if (ReadRes == 0) {
        mStatus = ConnectionBuffer::STATUS_COMPLETED;
    }
    else {
        if (mReadBuf.length() + (uint32_t) ReadRes > HTTP_RESPONSE_SIZE) {
            Reset();
            NAIPAX_THROW(Error::DATA_OUT_OF_RANGE);
        }

        // Winnow out unwanted response codes.
        if (mGotStatusCode != true) {
            uint32_t Code = HttpResponse::GetStatusCode(Buf, ReadRes);
            if (mpAcceptedStatus->find(Code) == mpAcceptedStatus->end()) {
                Reset();
                NAIPAX_THROW(Error::DATA_NOT_ACCEPTED);
            }
            mGotStatusCode = true;
        }
        mReadBuf += string(Buf, (uint32_t) ReadRes);
    }
}

void ConnectionBuffer::Connect() {
    mFd = socket(PF_INET, SOCK_STREAM, 0);
    if (mFd < 0) {
        Reset();
        NAIPAX_THROW(Error::CANNOT_CONNECT);
    }

    // Getting the previous file status flag is not necessary here.
    if (fcntl(mFd, F_SETFL, O_NONBLOCK) < 0) {
        Reset();
        NAIPAX_THROW(Error::CANNOT_SET_NONBLOCK);
    }

    struct sockaddr_in ServerAddr;
    memset(&ServerAddr, 0, sizeof(ServerAddr));
    ServerAddr.sin_family = AF_INET;
    memcpy(&(ServerAddr.sin_addr.s_addr), &mAresData.mIpV4Addr,
           sizeof(struct in_addr));
    ServerAddr.sin_port = htons(atol(mUrl.Port().c_str()));

    int32_t Ret = connect(mFd, (struct sockaddr *) &ServerAddr,
                          sizeof(ServerAddr));
    if (Ret == -1 && errno != EINPROGRESS) {
        Reset();
        NAIPAX_THROW(Error::CANNOT_CONNECT);
    }

    mStatus = ConnectionBuffer::STATUS_CONNECTING;
}

void ConnectionBuffer::BuildRequest() {
    int32_t Res = 0;
    socklen_t Size = sizeof(Res);
    if (getsockopt(mFd, SOL_SOCKET, SO_ERROR, &Res, &Size) != 0) {
        Reset();
        NAIPAX_THROW(Error::CANNOT_CONNECT);
    }
    if (Res) {
        Reset();
        NAIPAX_THROW(Error::CANNOT_CONNECT);
    }

    try {
        HttpRequest Req;
        if (mpUserAgent != NULL) {
            Req.UserAgent() = *mpUserAgent;
        }
        Req.Build(mUrl.Host().c_str(),
                  ("/" + mUrl.Path() + (mUrl.Query().length() ?
                                        "?" + mUrl.Query() : "")).c_str());
        // Don't pass zero buffer into write()
        if (Req.Headers().length() == 0) {
            NAIPAX_THROW(Error::CANNOT_BUILD_DATA);
        }
        mWriteBuf = Req.Headers();
        mStatus = ConnectionBuffer::STATUS_CONNECTED;
    }
    catch (Naipax::Error &E) {
        Reset();
        NAIPAX_THROW(Error::CANNOT_BUILD_DATA);
    }
}

Connection::Connection(uint32_t MaxConnCount) : mpResponseHandler(NULL),
                                                _mCompletedCount(0),
                                                _mMaxConnCount(MaxConnCount) {
    _mAcceptedStatus.insert(200);
    try {
        _mpConnBuf = new ConnectionBuffer[MaxConnCount];
        _mpPollFds = new struct pollfd[MaxConnCount];
    }
    catch (std::bad_alloc &E) {
        NAIPAX_THROW(Error::CANNOT_INIT_CONNECTION);
    }
}

Connection::~Connection() {
    delete[] _mpPollFds;
    delete[] _mpConnBuf;
}

uint32_t Connection::CurConnCount() {
    uint32_t Count = 0;
    for (uint32_t i = 0; i < _mMaxConnCount; ++i) {
        if (_mpConnBuf[i].mStatus != ConnectionBuffer::STATUS_INIT) {
            ++Count;
        }
    }
    return Count;
}

void Connection::Add(const std::string &rUrl, time_t Timeout) {
    if (_mAcceptedStatus.empty()) {
        NAIPAX_THROW(Error::DATA_NOT_ACCEPTED);
    }
    if (rUrl.length() == 0 || Timeout == 0) {
        return;
    }
    
    // Find an empty slot and find the maximum file descriptor for
    // now.
    uint32_t Slot = _mMaxConnCount;
    for (uint32_t i = 0; i < _mMaxConnCount; ++i) {
        ConnectionBuffer &rCurConnBuf = _mpConnBuf[i];
        if (rCurConnBuf.mStatus == ConnectionBuffer::STATUS_INIT) {
            Slot = i;
            break;
        }
    }
    if (Slot == _mMaxConnCount) {
        NAIPAX_THROW(Error::BUFFER_IS_FULL);
    }

    ConnectionBuffer &rConnBuf = _mpConnBuf[Slot];

    rConnBuf.mStatus = ConnectionBuffer::STATUS_UNRESOLVED;

    // Parse url and normalize url
    rConnBuf.mUrl.Set(rUrl);
    if (rConnBuf.mUrl.IsValid() != true) {
        rConnBuf.Reset();
        NAIPAX_THROW(Error::INVALID_ARGUMENT);
    }
    rConnBuf.mUrl.Normalize();
    if (rConnBuf.mUrl.AsString().length() == 0
        || rConnBuf.mUrl.Host().length() == 0) {
        rConnBuf.Reset();
        NAIPAX_THROW(Error::INVALID_ARGUMENT);
    }
    rConnBuf.mExpireTime = time(NULL) + Timeout;
    rConnBuf.mpUserAgent = &(this->_mUserAgent);
    rConnBuf.mpAcceptedStatus = &(this->_mAcceptedStatus);
}

void Connection::OneNbLoop() {
    // Check domain name resolution
    if (_mAres.HasUnresolved() == true) {
        vector<AresData> DnDataVector;
        if (_mAres.HasResult(DnDataVector) == true) {
            for (vector<AresData>::iterator DataIter = DnDataVector.begin();
                 DataIter != DnDataVector.end(); ++DataIter) {
                if (DataIter->mIsIpV6 == false) {
                    _mpConnBuf[DataIter->mIndex].mStatus
                        = ConnectionBuffer::STATUS_RESOLVED;
                    _mpConnBuf[DataIter->mIndex].mAresData = *DataIter;
                }
                else {
                    // IPv6 is not handled for now.
                    // Reset buffer status
                    _mpConnBuf[DataIter->mIndex].Reset();
                }
            }
        }
    }
    
    // The following actions are assumed to happen within one second.
    // Thus time() is called only once.
    time_t CurTime = time(NULL);

    // (poll fd => the index of connection buffer)
    map<nfds_t, uint32_t> PollFdMap;
    nfds_t PollFdCount = 0;
    
    for (uint32_t i = 0; i < _mMaxConnCount; ++i) {
        ConnectionBuffer &rCurConnBuf = _mpConnBuf[i];
        uint32_t &rStatus = rCurConnBuf.mStatus;

        if (rStatus & ConnectionBuffer::STATUS_INIT) {
            continue;
        }

        if (CurTime > rCurConnBuf.mExpireTime) {
            rCurConnBuf.Reset();
            continue;
        }

        if (rStatus & ConnectionBuffer::STATUS_RESOLVING) {
            continue;
        }

        try {
            if (rStatus & ConnectionBuffer::STATUS_UNRESOLVED) {
                AresData DnData(rCurConnBuf.mUrl.Host().c_str(),
                                rCurConnBuf.mUrl.Host().length());
                DnData.mIndex = i;
                if (_mAres.Submit(DnData) == -1) {
                    continue;
                }
                rStatus = ConnectionBuffer::STATUS_RESOLVING;
            }
            else if (rStatus & ConnectionBuffer::STATUS_RESOLVED) {
                rCurConnBuf.Connect();
            }
            else if (rStatus & ConnectionBuffer::STATUS_CONNECTING) {
                rCurConnBuf.BuildRequest();
            }
            else if (rStatus & (ConnectionBuffer::STATUS_CONNECTED
                                | ConnectionBuffer::STATUS_WRITTEN)) {
                _mpPollFds[PollFdCount].fd = rCurConnBuf.mFd;
                _mpPollFds[PollFdCount].events
                    = (rStatus & ConnectionBuffer::STATUS_CONNECTED ?
                       POLLOUT : POLLIN);
                PollFdMap[PollFdCount] = i;
                ++PollFdCount;
            }
            else if (rStatus & ConnectionBuffer::STATUS_COMPLETED) {
                ++_mCompletedCount;
                if (mpResponseHandler != NULL) {
                    try {
                        mpResponseHandler(*this,
                                          rCurConnBuf.mUrl.AsString(),
                                          rCurConnBuf.mReadBuf);
                    }
                    catch (Naipax::Error &E) {
                        cerr << E.What() << endl;
                    }
                    catch (...) {
                        // ResponseHandler should not break crawling
                        // process
                    }
                }
                rCurConnBuf.Reset();
            }
        }
        catch (Naipax::Error &E) {
            rCurConnBuf.Reset();
        }
    }

    // Read/Write happens in the following code.
    if (PollFdCount == 0) {
        return;
    }

    int32_t PollRet = poll(_mpPollFds, PollFdCount, 0);
    if (PollRet == 0) {
        return;
    }
    else if (PollRet < 0) {
        cerr << "Poll() error: " << strerror(errno)
             << " in function " << __FUNCTION__ << endl;
        return;
    }

    for (nfds_t i = 0; i < PollFdCount; ++i) {
        struct pollfd &rPollFd = _mpPollFds[i];
        ConnectionBuffer &rCurConnBuf = _mpConnBuf[PollFdMap[i]];
        uint32_t &rStatus = rCurConnBuf.mStatus;
        
        if (rStatus & (ConnectionBuffer::STATUS_CONNECTED
                       | ConnectionBuffer::STATUS_WRITTEN)) {
            try {
                if (rPollFd.revents & POLLIN) {
                    rCurConnBuf.ReadData();
                }
                else if (rPollFd.revents & POLLOUT) {
                    rCurConnBuf.WriteData();
                }
                else if (rPollFd.revents & (POLLERR | POLLHUP | POLLNVAL)) {
                    rCurConnBuf.Reset();
                }
            }
            catch (Naipax::Error &E) {
                rCurConnBuf.Reset();
            }
        }
    }
}

